"""
===============================
@Author     : Zuo WenTao
@Time       : 2023/11/10 13:22
@Description:
@Change     :
@File       : embedding_resquest.py
===============================
"""
import random
import requests
import time

from ftlangdetect import detect
from concurrent.futures.thread import ThreadPoolExecutor
from loguru import logger


class MyRetry:
    def __init__(self,
                 times=1,
                 return_obj=None,
                 return_msg=None,
                 msg_return_params=False,
                 return_err_cause=True,
                 show_process_err=True,
                 sleep_time=0
                 ):
        """
        重试模块
        :param times: 重试次数
        :param return_obj: 重试失败会返回的对象，如果参数为None则不返回直接raise，否则返回除None外指定对象
        :param return_msg: 自定义重试失败后打印的消息
        :param msg_return_params: 是否要打印传入的参数
        :param return_err_cause: 是否打印报错原因
        :param show_process_err: 是否打印重试过程中的提示
        :param sleep_time: 暂停时间
        """
        self.times = times
        self.return_obj = return_obj
        self.return_msg = return_msg
        self.msg_return_params = msg_return_params
        self.return_err_cause = return_err_cause
        self.show_process_err = show_process_err
        self.sleep_time = sleep_time

    def __call__(self, fun):
        retry_times = self.times
        return_obj = self.return_obj
        return_msg = self.return_msg
        msg_return_params = self.msg_return_params
        return_err_cause = self.return_err_cause
        show_process_err = self.show_process_err
        sleep_time = self.sleep_time

        def retry_handle(*args, **kwargs):
            msg = f'函数: {fun.__name__}   重试达到最大次数{retry_times}次'
            cause = ''
            for times in range(retry_times):
                try:
                    fun_res = fun(*args, **kwargs)
                except Exception as e:
                    cause = e
                    if times == retry_times - 1:
                        if return_obj is not None and return_msg is None:
                            msg += f' 返回自定义对象: {return_obj} '
                            return return_obj
                        elif return_obj is None and return_msg is not None:
                            msg += f' 消息: {return_msg} '
                        elif return_obj is not None and return_msg is not None:
                            msg += f' 返回自定义对象: {return_obj}  消息: {return_msg} '
                            return return_obj
                        else:
                            msg += f' 重试达到最大次数({retry_times}次 '
                            raise
                    else:
                        if show_process_err:
                            # logger.error(f'函数: {fun.__name__} 发生异常:{e}，第{times + 1}次重试:')
                            # logger.debug(*args)
                            pass
                        if sleep_time:
                            time.sleep(sleep_time)
                else:
                    return fun_res
                finally:
                    if times == retry_times - 1 and msg:
                        if msg_return_params:
                            msg += f' 参数: {args} {kwargs} '
                        if return_err_cause:
                            msg += f'原因: {cause}'
                        # logger.warning(msg)

        return retry_handle

ch_server_list = [
    "http://192.168.200.39:32003/embedding",
    "http://192.168.200.39:32004/embedding",
    "http://192.168.200.39:32005/embedding",
    "http://192.168.200.39:32006/embedding",
    "http://192.168.200.39:32010/embedding",
    "http://192.168.200.39:32011/embedding",
]

en_server_list = [
    "http://192.168.200.39:31003/embedding",
    "http://192.168.200.39:31004/embedding",
    "http://192.168.200.39:31005/embedding",
    "http://192.168.200.39:31006/embedding",
    "http://192.168.200.39:31010/embedding",
    "http://192.168.200.39:31011/embedding",
]

class GetEmbeddingRequest:
    _instance = None

    def __new__(cls, *args, **kwargs):
        if GetEmbeddingRequest._instance is None:
            GetEmbeddingRequest._instance = super().__new__(cls)
        return GetEmbeddingRequest._instance

    @staticmethod
    @MyRetry(times=3, return_msg="向量转换失败")
    def get_embedding(text, server):
        data = {'text': text}
        respon = requests.post(server, json=data, timeout=5)
        assert respon.status_code == 200, f'向量接口请求失败'
        res_data = respon.json()
        assert res_data['code'] == 200, res_data['msg']
        result = respon.json()["data"]
        assert result["embdding"], f'返回结果错误：{result}'
        if result["embdding"] is None:
            raise Exception(f'返回结果错误：{result}')
        return result["embdding"]


class GetEmbeddingManager:
    _instance = None

    def __init__(self):
        pass

    def __new__(cls, *args, **kwargs):
        if GetEmbeddingManager._instance is None:
            GetEmbeddingManager._instance = super().__new__(cls)
        return GetEmbeddingManager._instance

    def get_language(self, text):
        return detect(text=text, low_memory=False)['lang']

    def manager(self, text):
        """
        检测语言，选择对应语言的向量接口
        ch_server_list: 中文向量服务列表
        en_server_list: 英文向量服务列表
        :param text: 输入文本
        :return:
        """
        lang = self.get_language(text)
        if lang == 'zh':
            return GetEmbeddingRequest.get_embedding(text, random.choice(ch_server_list))
        else:
            return GetEmbeddingRequest.get_embedding(text, random.choice(en_server_list))


if __name__ == '__main__':
    # ceshi_list = [
    #     "http://192.168.200.39:32001/embedding"
    # ]
    # data = "张三"
    # manager = GetEmbeddingManager().manager(data)
    # print(manager)
    print(GetEmbeddingRequest().get_embedding("张三", "http://192.168.66.95:8891/embedding"))
